package com.jq.jqaicode.ai.workflow.ruotes;

import cn.hutool.core.util.StrUtil;
import com.jq.jqaicode.mapper.AiAppMapper;
import com.jq.jqaicode.model.entity.AiApp;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.response.ChatResponse;
import org.bsc.langgraph4j.action.EdgeAction;
import org.bsc.langgraph4j.prebuilt.MessagesState;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;

/**
 * AI路由
 */
@Component
public class CodeTypeRoute implements EdgeAction<MessagesState<ChatMessage>> {

    @Autowired
    private ChatModel chatModel;
    @Autowired
    AiAppMapper aiAppMapper;

    @Override
    public String apply(MessagesState<ChatMessage> chatMessageMessagesState) throws Exception {
        UserMessage chatMessage = (UserMessage)chatMessageMessagesState.lastMessage()
                .orElseThrow(() -> (new IllegalStateException("last message not found!")));
        //更新应用类型
        AiApp app = aiAppMapper.selectById(chatMessage.name());
        String codeGenType=app.getCodeGenType();
        //如果项目类型是空，使用大模型LLM选择类型
        if(StrUtil.isBlank(codeGenType)){
            SystemMessage systemMessage= SystemMessage.from("你是一位资深的 Web 前端开发专家，精通 HTML和VUE的前端项目设计和构建，\n" +
                    "你的任务是根据用户提供的网站描述，判断用户返回要生成什么类型的项目，有两种类型区分：html和vue。\n" +
                    "\n" +
                    "如果用户描述是复杂的的项目，包含vue、node、element-ui等关键词就返回结果：vue"+
                    "如果用户描述是简单的项目，包含原生html等关键词就返回结果：html\n"+
                    "或者你分不出的，默认返回结果：html\n"+
                    "\n" +
                    "强制约束输出格式: html或vue\n");
            List<ChatMessage> messages = new ArrayList<>();
            messages.add(systemMessage);
            messages.add(chatMessage);
            ChatRequest request = ChatRequest.builder()
                    .messages(messages)
                    .build();

            ChatResponse chat = chatModel.chat(request);
            if("vue".equals(chat.aiMessage().text())){
                codeGenType="vue";
            }else{
                //html类型
                codeGenType="html";
            }
            app.setCodeGenType(codeGenType);
            aiAppMapper.updateById(app);
        }
        return codeGenType;
    }

}
